{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "46d6c164544f"
      },
      "outputs": [],
      "source": [
        "# Copyright 2020 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rEJBSTyZIrIb"
      },
      "source": [
        "# Training PyTorch Model on Google Cloud AI Platform Training \n",
        "## Fine Tuning Pretrained [BERT](https://huggingface.co/bert-base-cased) Model for Sentiment Classification Task "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X4cRE8IbIrIV"
      },
      "source": [
        "# Overview\n",
        "\n",
        "This example is inspired from Token-Classification [notebook](https://github.com/huggingface/notebooks/blob/master/examples/token_classification.ipynb) and [run_glue.py](https://github.com/huggingface/transformers/blob/v2.5.0/examples/run_glue.py) from HuggingFace 🤗. \n",
        "We will be fine-tuning bert-base-cased (pre-trained) model.\n",
        "You can find the details about this model at [🤗  Hub](https://huggingface.co/bert-base-cased).\n",
        "\n",
        "For more notebooks of the state of the art PyTorch/Tensorflow/JAX you can explore [🤗 Notebooks](https://huggingface.co/transformers/notebooks.html).\n",
        "\n",
        "### Dataset\n",
        "\n",
        "We will be using IMDB moview review dataset from Huggingface Datasets.\n",
        "\n",
        "### Objective\n",
        "\n",
        "Get familiar with PyTorch on Cloud AI Platform notebooks instances.\n",
        "\n",
        "### Costs \n",
        "\n",
        "This tutorial uses billable components of Google Cloud Platform (GCP):\n",
        "\n",
        "* Cloud AI Platform Notebook\n",
        "* Cloud AI Platform Training\n",
        "\n",
        "Learn about [Cloud AI Platform\n",
        "pricing](https://cloud.google.com/ml-engine/docs/pricing) and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cda8ab8b7a04"
      },
      "source": [
        "## Setting up Notebook Environment\n",
        "\n",
        "This notebook assumes [PyTorch 1.7 DLVM](https://cloud.google.com/ai-platform/notebooks/docs/images) development environment. You can create a Notebook instance using [Google Cloud Console](https://cloud.google.com/ai-platform/notebooks/docs/create-new) or [`gcloud` command](https://cloud.google.com/sdk/gcloud/reference/notebooks/instances/create).\n",
        "\n",
        "```\n",
        "gcloud notebooks instances create example-instance \\\n",
        "    --vm-image-project=deeplearning-platform-release \\\n",
        "    --vm-image-family=pytorch-1-7-cu110-notebooks \\\n",
        "    --machine-type=n1-standard-4 \\\n",
        "    --location=us-central1-a \\\n",
        "    --boot-disk-size=100 \\\n",
        "    --accelerator-core-count=1 \\\n",
        "    --accelerator-type=NVIDIA_TESLA_T4 \\\n",
        "    --install-gpu-driver \\\n",
        "    --network=default\n",
        "```\n",
        "\n",
        "**NOTE:** You must have GPU quota before you can create instances with GPUs. Check the [quotas](https://console.cloud.google.com/iam-admin/quotas) page to ensure that you have enough GPUs available in your project. If GPUs are not listed on the quotas page or you require additional GPU quota, [request a quota increase](https://cloud.google.com/compute/quotas#requesting_additional_quota). Free Trial accounts do not receive GPU quota by default."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9a864c22307c"
      },
      "source": [
        "### Python Dependencies\n",
        "\n",
        "Python dependencies required for this notebook are [Transformers](https://pypi.org/project/transformers/) and [Datasets](https://pypi.org/project/datasets/) and will be installed in the notebook itself."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2974254ea6be"
      },
      "outputs": [],
      "source": [
        "!pip -q install torch==1.7\n",
        "!pip -q install transformers\n",
        "!pip -q install datasets\n",
        "!pip -q install tqdm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f361541eff05"
      },
      "source": [
        "### Restart the Kernel\n",
        "\n",
        "Once you've installed the {packages}, you need to restart the notebook kernel so it can find the packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2d77a223d63d"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "276e063f41eb"
      },
      "source": [
        "### Python imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f656723ac8ed"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from datasets import load_dataset\n",
        "from transformers import (AutoModelForSequenceClassification, AutoTokenizer,\n",
        "                          EvalPrediction, Trainer, TrainingArguments,\n",
        "                          default_data_collator)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "whPRbBNbIrIl"
      },
      "source": [
        "## Loading the dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W7QYTpxXIrIl"
      },
      "source": [
        "We use the [🤗 Datasets](https://github.com/huggingface/datasets) library to download the data and get the metric we need to use for evaluation (to compare our model to the benchmark). This can be easily done with the functions `load_dataset` and `load_metric`.\n",
        "\n",
        "For this example we will use IMDB movie review dataset for sentiment classification task."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IreSlFmlIrIm"
      },
      "outputs": [],
      "source": [
        "datasets = load_dataset(\"imdb\")\n",
        "batch_size = 16\n",
        "max_seq_length = 128\n",
        "model_name_or_path = \"bert-base-cased\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9114cb03f4cf"
      },
      "outputs": [],
      "source": [
        "datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RzfPtOMoIrIu"
      },
      "source": [
        "The `datasets` object itself is [`DatasetDict`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasetdict), which contains one key for the training, validation and test set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GWiVUF0jIrIv"
      },
      "outputs": [],
      "source": [
        "print(\n",
        "    \"Total # of rows in training dataset {} and size {:5.2f} MB\".format(\n",
        "        datasets[\"train\"].shape[0], datasets[\"train\"].size_in_bytes / (1024 * 1024)\n",
        "    )\n",
        ")\n",
        "print(\n",
        "    \"Total # of rows in test dataset {} and size {:5.2f} MB\".format(\n",
        "        datasets[\"test\"].shape[0], datasets[\"test\"].size_in_bytes / (1024 * 1024)\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u3EtYfeHIrIz"
      },
      "source": [
        "To access an actual element, you need to select a split first, then give an index:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X6HrpprwIrIz"
      },
      "outputs": [],
      "source": [
        "datasets[\"train\"][0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dQPMGfEFLhFd"
      },
      "source": [
        "Using the `unique` method to extract label list. This will allow us to experiment with other datasets without hard-coding labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TYw3eIHfLhFd"
      },
      "outputs": [],
      "source": [
        "label_list = datasets[\"train\"].unique(\"label\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WHUmphG3IrI3"
      },
      "source": [
        "To get a sense of what the data looks like, the following function will show some examples picked randomly in the dataset (automatically decoding the labels in passing)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i3j8APAoIrI3"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "\n",
        "import pandas as pd\n",
        "from datasets import ClassLabel, Sequence\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "\n",
        "def show_random_elements(dataset, num_examples=2):\n",
        "    assert num_examples <= len(\n",
        "        dataset\n",
        "    ), \"Can't pick more elements than there are in the dataset.\"\n",
        "    picks = []\n",
        "    for _ in range(num_examples):\n",
        "        pick = random.randint(0, len(dataset) - 1)\n",
        "        while pick in picks:\n",
        "            pick = random.randint(0, len(dataset) - 1)\n",
        "        picks.append(pick)\n",
        "\n",
        "    df = pd.DataFrame(dataset[picks])\n",
        "    for column, typ in dataset.features.items():\n",
        "        if isinstance(typ, ClassLabel):\n",
        "            df[column] = df[column].transform(lambda i: typ.names[i])\n",
        "        elif isinstance(typ, Sequence) and isinstance(typ.feature, ClassLabel):\n",
        "            df[column] = df[column].transform(\n",
        "                lambda x: [typ.feature.names[i] for i in x]\n",
        "            )\n",
        "    display(HTML(df.to_html()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SZy5tRB_IrI7"
      },
      "outputs": [],
      "source": [
        "show_random_elements(datasets[\"train\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n9qywopnIrJH"
      },
      "source": [
        "## Preprocessing the data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YVx71GdAIrJH"
      },
      "source": [
        "Before we can feed those texts to our model, we need to preprocess them. This is done by a 🤗 Transformers `Tokenizer` which will (as the name indicates) tokenize the inputs (including converting the tokens to their corresponding IDs in the pretrained vocabulary) and put it in a format the model expects, as well as generate the other inputs that model requires.\n",
        "\n",
        "To do all of this, we instantiate our tokenizer with the `AutoTokenizer.from_pretrained` method, which will ensure:\n",
        "\n",
        "- we get a tokenizer that corresponds to the model architecture we want to use,\n",
        "- we download the vocabulary used when pretraining this specific checkpoint.\n",
        "\n",
        "That vocabulary will be cached, so it's not downloaded again the next time we run the cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eXNLu_-nIrJI"
      },
      "outputs": [],
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(\n",
        "    model_name_or_path,\n",
        "    use_fast=True,\n",
        ")\n",
        "# 'use_fast' ensure that we use fast tokenizers (backed by Rust) from the 🤗 Tokenizers library."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "muc7yghlLhFf"
      },
      "source": [
        "You can check which type of models have a fast tokenizer available and which don't on the [big table of models](https://huggingface.co/transformers/index.html#bigtable)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rowT4iCLIrJK"
      },
      "source": [
        "You can directly call this tokenizer on one sentence:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a5hBlsrHIrJL"
      },
      "outputs": [],
      "source": [
        "tokenizer(\"Hello, this is one sentence!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qeT4rxOGLhFf"
      },
      "source": [
        "Depending on the model you selected, you will see different keys in the dictionary returned by the cell above. They don't matter much for what we're doing here (just know they are required by the model we will instantiate later), you can learn more about them in [this tutorial](https://huggingface.co/transformers/preprocessing.html) if you're interested.\n",
        "\n",
        "Note: If, as is the case here, your inputs have already been split into words, you should pass the list of words to your tokenzier with the argument `is_split_into_words=True`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MPunumj-LhFg"
      },
      "outputs": [],
      "source": [
        "example = datasets[\"train\"][4]\n",
        "print(example)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e3-m0sGHLhFg"
      },
      "outputs": [],
      "source": [
        "tokenizer(\n",
        "    [\"Hello\", \",\", \"this\", \"is\", \"one\", \"sentence\", \"split\", \"into\", \"words\", \".\"],\n",
        "    is_split_into_words=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G73hFICnLhFg"
      },
      "source": [
        "Note that transformers are often pretrained with subword tokenizers, meaning that even if your inputs have been split into words already, each of those words could be split again by the tokenizer. Let's look at an example of that:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d8c6f3c11353"
      },
      "outputs": [],
      "source": [
        "# Dataset loading repeated here to make this cell idempotent\n",
        "# Since we are over-writing datasets variable\n",
        "datasets = load_dataset(\"imdb\")\n",
        "\n",
        "# TEMP: We can extract this automatically but Unique method of the dataset\n",
        "# is not reporting the label -1 which shows up in the pre-processing\n",
        "# Hence the additional -1 term in the dictionary\n",
        "label_to_id = {1: 1, 0: 0, -1: 0}\n",
        "\n",
        "\n",
        "def preprocess_function(examples):\n",
        "    # Tokenize the texts\n",
        "    args = (examples[\"text\"],)\n",
        "    result = tokenizer(\n",
        "        *args, padding=\"max_length\", max_length=max_seq_length, truncation=True\n",
        "    )\n",
        "\n",
        "    # Map labels to IDs (not necessary for GLUE tasks)\n",
        "    if label_to_id is not None and \"label\" in examples:\n",
        "        result[\"label\"] = [label_to_id[example] for example in examples[\"label\"]]\n",
        "\n",
        "    return result\n",
        "\n",
        "\n",
        "datasets = datasets.map(preprocess_function, batched=True, load_from_cache_file=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "545PP3o8IrJV"
      },
      "source": [
        "## Fine Tuning the Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FBiW8UpKIrJW"
      },
      "source": [
        "Now that our data is ready, we can download the pretrained model and fine-tune it. Since all our tasks are about token classification, we use the `AutoModelForSequenceClassification` class. Like with the tokenizer, the `from_pretrained` method will download and cache the model for us. The only thing we have to specify is the number of labels for our problem (which we can get from the features, as seen before):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TlqNaB8jIrJW"
      },
      "outputs": [],
      "source": [
        "model = AutoModelForSequenceClassification.from_pretrained(\n",
        "    model_name_or_path, num_labels=len(label_list)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CczA5lJlIrJX"
      },
      "source": [
        "The warning is telling us we are throwing away some weights (the `vocab_transform` and `vocab_layer_norm` layers) and randomly initializing some other (the `pre_classifier` and `classifier` layers). This is absolutely normal in this case, because we are removing the head used to pretrain the model on a masked language modeling objective and replacing it with a new head for which we don't have pretrained weights, so the library warns us we should fine-tune this model before using it for inference, which is exactly what we are going to do."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_N8urzhyIrJY"
      },
      "source": [
        "To instantiate a `Trainer`, we will need to define three more things. The most important is the [`TrainingArguments`](https://huggingface.co/transformers/main_classes/trainer.html#transformers.TrainingArguments), which is a class that contains all the attributes to customize the training. It requires one folder name, which will be used to save the checkpoints of the model, and all other arguments are optional:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bliy8zgjIrJY"
      },
      "outputs": [],
      "source": [
        "args = TrainingArguments(\n",
        "    evaluation_strategy=\"epoch\",\n",
        "    learning_rate=2e-5,\n",
        "    per_device_train_batch_size=batch_size,\n",
        "    per_device_eval_batch_size=batch_size,\n",
        "    num_train_epochs=1,\n",
        "    weight_decay=0.01,\n",
        "    output_dir=\"/tmp/cls\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "km3pGVdTIrJc"
      },
      "source": [
        "Here we set the evaluation to be done at the end of each epoch, tweak the learning rate, use the `batch_size` defined at the top of the notebook and customize the number of epochs for training, as well as the weight decay."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AqynGv2DLhFk"
      },
      "source": [
        "The last thing to define for our `Trainer` is how to compute the metrics from the predictions. You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a predictions and label_ids field) and has to return a dictionary string to float."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V233MNZgLhFk"
      },
      "outputs": [],
      "source": [
        "def compute_metrics(p: EvalPrediction):\n",
        "    preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions\n",
        "    preds = np.argmax(preds, axis=1)\n",
        "    return {\"accuracy\": (preds == p.label_ids).astype(np.float32).mean().item()}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aehFUe14LhFl"
      },
      "source": [
        "Now we Create the Trainer object and we are almost ready to train."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "imY1oC3SIrJf"
      },
      "outputs": [],
      "source": [
        "trainer = Trainer(\n",
        "    model,\n",
        "    args,\n",
        "    train_dataset=datasets[\"train\"],\n",
        "    eval_dataset=datasets[\"test\"],\n",
        "    data_collator=default_data_collator,\n",
        "    tokenizer=tokenizer,\n",
        "    compute_metrics=compute_metrics,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CdzABDVcIrJg"
      },
      "source": [
        "We can now finetune our model by just calling the `train` method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wZHmdWPELhFl"
      },
      "outputs": [],
      "source": [
        "trainer.train()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "58ab9c74b4f8"
      },
      "outputs": [],
      "source": [
        "trainer.save_model(\"./finetuned-bert-classifier\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CKASz-2vIrJi"
      },
      "source": [
        "The `evaluate` method allows you to evaluate again on the evaluation dataset or on another dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UOUcBkX8IrJi"
      },
      "outputs": [],
      "source": [
        "trainer.evaluate()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gX97-EH5LhFm"
      },
      "source": [
        "To get the precision/recall/f1 computed for each category now that we have finished training, we can apply the same function as before on the result of the `predict` method:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5cae4fdf26f2"
      },
      "source": [
        "## Running Predictions with Sample Examples"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "21ab30d5518c"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import AutoModelForSequenceClassification, AutoTokenizer\n",
        "\n",
        "model_name_or_path = \"bert-base-cased\"\n",
        "label_text = {0: \"Negative\", 1: \"Positive\"}\n",
        "saved_model_path = \"./finetuned-bert-classifier\"\n",
        "\n",
        "\n",
        "def predict(input_text, saved_model_path):\n",
        "    # initialize tokenizer\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)\n",
        "\n",
        "    # preprocess and encode input text\n",
        "    predict_input = tokenizer.encode(\n",
        "        review_text, truncation=True, max_length=128, return_tensors=\"pt\"\n",
        "    )\n",
        "\n",
        "    # load trained model\n",
        "    loaded_model = AutoModelForSequenceClassification.from_pretrained(saved_model_path)\n",
        "\n",
        "    # get predictions\n",
        "    output = loaded_model(predict_input)\n",
        "\n",
        "    # return labels\n",
        "    label_id = torch.argmax(*output.to_tuple(), dim=1)\n",
        "\n",
        "    print(f\"Review text: {review_text}\")\n",
        "    print(f\"Sentiment : {label_text[label_id.item()]}\\n\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e594d72fa8e4"
      },
      "outputs": [],
      "source": [
        "# example #1\n",
        "review_text = (\n",
        "    \"\"\"Jaw dropping visual affects and action! One of the best I have seen to date.\"\"\"\n",
        ")\n",
        "predict_input = predict(review_text, saved_model_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dd393f0d38c1"
      },
      "outputs": [],
      "source": [
        "# example #2\n",
        "review_text = \"\"\"Take away the CGI and the A-list cast and you end up with film with less punch.\"\"\"\n",
        "predict_input = predict(review_text, saved_model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "46a8d95366ee"
      },
      "source": [
        "## Run Training Job on Cloud AI Platform (CAIP)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e97e8af2376c"
      },
      "source": [
        "You can do local experimentation on your AI Platform Notebooks instance. However, for larger datasets or models often a vertically scaled compute or horizontally distributed training is required. The most cost effective way to perform this task is [Cloud AI Platform Training](https://cloud.google.com/ai-platform/training/docs) Service. AI Platform Training takes care of creating designated compute resources, performs the training task and ensures deletion of compute resources once the training job is finished.\n",
        "\n",
        "In this part of the notebook, we will show you scaling your training job by packaging the code and submitting the training job to AI Platform Training."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "585ba542aa40"
      },
      "source": [
        "### Packaging the Training Application\n",
        "\n",
        "Before runnning the training application with AI Platform Training, training application code and any dependencies must be uploaded into a Cloud Storage bucket that your Google Cloud project can access. This sections shows how to package and stage your application in the cloud.\n",
        "\n",
        "There are two ways to package your application and dependencies and run on AI Platform Training:\n",
        "\n",
        "1. Package [application and Python dependencies](https://cloud.google.com/ai-platform/training/docs/packaging-trainer#working_with_dependencies) manually using setup tools\n",
        "2. Use [custom containers](https://cloud.google.com/ai-platform/training/docs/custom-containers-training) to package dependencies using Docker containers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4b8b5333b577"
      },
      "source": [
        "#### Recommended Training Application Structure\n",
        "\n",
        "You can structure your training application in any way you like. However, the [following structure](https://cloud.google.com/ai-platform/training/docs/packaging-trainer#project-structure) is commonly used in AI Platform Training samples, and having your project's organization be similar to the samples can make it easier for you to follow the samples.\n",
        "\n",
        "We have two directories `python_package` and `custom_container` showing both the packaging approaches. `README.md` files inside each directory has details on the directory structure and instructions on howw to run application locally and on the cloud.\n",
        "\n",
        "```\n",
        ".\n",
        "├── custom_container\n",
        "│   ├── Dockerfile\n",
        "│   ├── README.md\n",
        "│   ├── scripts\n",
        "│   │   ├── train-cloud.sh\n",
        "│   │   └── train-local.sh\n",
        "│   └── trainer -> ../python_package/trainer/\n",
        "├── python_package\n",
        "│   ├── README.md\n",
        "│   ├── scripts\n",
        "│   │   ├── train-cloud.sh\n",
        "│   │   └── train-local.sh\n",
        "│   ├── setup.py\n",
        "│   └── trainer\n",
        "│       ├── __init__.py\n",
        "│       ├── experiment.py\n",
        "│       ├── metadata.py\n",
        "│       ├── model.py\n",
        "│       ├── task.py\n",
        "│       └── utils.py\n",
        "└── pytorch-text-classification-caip-training.ipynb    --> This notebook\n",
        "```\n",
        "\n",
        "1. Main project directory contains your `setup.py` file or `Dockerfile` with the dependencies. \n",
        "2. Use a subdirectory named `trainer` to store your main application module and `scripts` to submit training jobs locally or cloud\n",
        "3. Inside `trainer` directory:\n",
        "    - `task.py` - Main application module 1) initialises and parse task arguments (hyper parameters), and 2) entry point to the trainer\n",
        "    - `model.py` -  Includes function to create model with a sequence classification head from a pretrained model.\n",
        "    - `experiment.py` - Runs the model training and evaluation experiment, and exports the final model.\n",
        "    - `metadata.py` - Defines metadata for classification task such as predefined model dataset name, target labels\n",
        "    - `utils.py` - Includes utility functions such as data input functions to read data, save model to GCS bucket\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f7466d414a0e"
      },
      "source": [
        "### Using Python Packaging to Build Manually"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1a675e6ee9dc"
      },
      "source": [
        "In this notebook, we are using Huggingface datasets and fine tuning a transformer model from Huggingface Transformers Library for sentiment analysis task. We will be adding standard Python dependencies - `transformers`, `datasets` and `tqdm` - in the `setup.py` file. The `find_packages()` function inside `setup.py` includes the `trainer` directory in the package as it contains `__init__.py` which tells [Python Setuptools](https://setuptools.readthedocs.io/en/latest/) to include all subdirectories of the parent directory as dependencies.\n",
        "\n",
        "```\n",
        "# ==========================================\n",
        "# contents of setup.py file\n",
        "# ==========================================\n",
        "\n",
        "from setuptools import find_packages\n",
        "from setuptools import setup\n",
        "\n",
        "REQUIRED_PACKAGES = [\n",
        "    'torch==1.7',\n",
        "    'transformers',\n",
        "    'datasets',\n",
        "    'tqdm'\n",
        "]\n",
        "\n",
        "setup(\n",
        "    name='trainer',\n",
        "    version='0.1',\n",
        "    install_requires=REQUIRED_PACKAGES,\n",
        "    packages=find_packages(),\n",
        "    include_package_data=True,\n",
        "    description='AI Platform | Training | PyTorch | Text Classification | Python Package'\n",
        ")\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "daf738249701"
      },
      "source": [
        "#### Running Training Job Locally"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dae35b9f526e"
      },
      "source": [
        "Before submitting the job to cloud, ensure the script runs locally. The script `./python_package/scripts/train-local.sh` runs training locally using `python -m trainer.task`\n",
        "\n",
        "```\n",
        "python -m trainer.task \\\n",
        "    --job-dir ${JOB_DIR} \\\n",
        "    --model-name=\"finetuned-bert-classifier\"\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c77fd0f20df0"
      },
      "outputs": [],
      "source": [
        "!cd python_package && ./scripts/train-local.sh"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4a1a4281ff7f"
      },
      "source": [
        "#### Running Training Job on Cloud AI Platform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3e3bd41fd9c3"
      },
      "source": [
        "You would submit the training job to Cloud AI Platform Training using `gcloud ai-platform jobs submit training`. `gcloud` command stages your training application on GCS bucket and submits the training job.\n",
        "\n",
        "\n",
        "```\n",
        "gcloud ai-platform jobs submit training ${JOB_NAME} \\\n",
        "    --region ${REGION} \\\n",
        "    --master-image-uri ${IMAGE_URI} \\\n",
        "    --scale-tier=CUSTOM \\\n",
        "    --master-machine-type=n1-standard-8 \\\n",
        "    --master-accelerator=type=nvidia-tesla-t4,count=2 \\\n",
        "    --job-dir ${JOB_DIR} \\\n",
        "    --module-name trainer.task \\\n",
        "    --package-path ${PACKAGE_PATH} \\\n",
        "    -- \\\n",
        "    --model-name=\"finetuned-bert-classifier\"\n",
        "```\n",
        "\n",
        "- Set the `--master-image-uri` flag to `gcr.io/cloud-aiplatform/training/pytorch-gpu.1-7` for training on pre-built PyTorch v1.7 image for GPU\n",
        "- Set the `--packages` flag to the path to your packaged application\n",
        "- Set the `--module-name` flag to the `trainer.task` which is the main module to start your application\n",
        "- Set the `--master-accelerator` and `--master-machine-type` flag to set the infrastructure to run the application. Refer [documentation](https://cloud.google.com/ai-platform/training/docs/machine-types) to set machine types and scaling tiers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e4519f1cfc75"
      },
      "outputs": [],
      "source": [
        "!cd python_package && ./scripts/train-cloud.sh"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c170d386492b"
      },
      "source": [
        "### Using Custom Containers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "035227b6e581"
      },
      "source": [
        "To create a training job with custom container, you have define a `Dockerfile` to install the dependencies required for the training job. Then, you build and test your Docker image locally to verify it before using it with AI Platform Training.\n",
        "\n",
        "\n",
        "```\n",
        "# ==========================================\n",
        "# contents of Dockerfile\n",
        "# ==========================================\n",
        "\n",
        "# Install pytorch\n",
        "FROM gcr.io/cloud-aiplatform/training/pytorch-gpu.1-7\n",
        "\n",
        "WORKDIR /root\n",
        "\n",
        "# Installs pandas, and google-cloud-storage.\n",
        "RUN pip install google-cloud-storage transformers datasets tqdm\n",
        "\n",
        "# Copies the trainer code to the docker image.\n",
        "COPY ./trainer/__init__.py ./trainer/__init__.py\n",
        "COPY ./trainer/experiment.py ./trainer/experiment.py\n",
        "COPY ./trainer/utils.py ./trainer/utils.py\n",
        "COPY ./trainer/metadata.py ./trainer/metadata.py\n",
        "COPY ./trainer/model.py ./trainer/model.py\n",
        "COPY ./trainer/task.py ./trainer/task.py\n",
        "\n",
        "# Set up the entry point to invoke the trainer.\n",
        "ENTRYPOINT [\"python\", \"-m\", \"trainer.task\"]\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e2be59c793d4"
      },
      "source": [
        "#### Running Training Job Locally with Custom Container\n",
        "\n",
        "Before submitting the job to cloud, ensure the script runs locally. The script `./python_package/scripts/train-local.sh` runs training locally using `python -m trainer.task`\n",
        "\n",
        "```\n",
        "# Build the docker image\n",
        "docker build -f Dockerfile -t ${IMAGE_URI} ./\n",
        "\n",
        "# Test your docker image locally\n",
        "echo \"Running the Docker Image\"\n",
        "docker run ${IMAGE_URI} \\\n",
        "    --job-dir ${JOB_DIR} \\\n",
        "    --model-name=\"finetuned-bert-classifier\"\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1f358291584e"
      },
      "outputs": [],
      "source": [
        "!cd custom_container && ./scripts/train-local.sh"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1873a991a40f"
      },
      "source": [
        "#### Running Training Job on Cloud AI Platform with Custom Container"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "64e706007344"
      },
      "source": [
        "Before submitting the training job, you need to push image to Google Cloud Container Registry and then submit the training job to Cloud AI Platform Training using `gcloud ai-platform jobs submit training`. \n",
        "\n",
        "\n",
        "```\n",
        "# Deploy the docker image to Cloud Container Registry\n",
        "docker push ${IMAGE_URI}\n",
        "\n",
        "# Submit the training job\n",
        "gcloud ai-platform jobs submit training ${JOB_NAME} \\\n",
        "    --region ${REGION} \\\n",
        "    --master-image-uri ${IMAGE_URI} \\\n",
        "    --scale-tier=CUSTOM \\\n",
        "    --master-machine-type=n1-standard-8 \\\n",
        "    --master-accelerator=type=nvidia-tesla-t4,count=2 \\\n",
        "    --job-dir ${JOB_DIR} \\\n",
        "    -- \\\n",
        "    --model-name=\"finetuned-bert-classifier\"\n",
        "```\n",
        "\n",
        "- Set the `--master-image-uri` flag to the custom container image pushed to Google Cloud Container Registry\n",
        "- Set the `--master-accelerator` and `--master-machine-type` flag to set the infrastructure to run the application. Refer [documentation](https://cloud.google.com/ai-platform/training/docs/machine-types) to set machine types and scaling tiers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "45c0f675d26a"
      },
      "outputs": [],
      "source": [
        "!cd custom_container && ./scripts/train-cloud.sh"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b329c52e304a"
      },
      "source": [
        "## Monitoring Training Job on Cloud AI Platform (CAIP)\n",
        "\n",
        "After you submit your job, you can monitor the job status using `gcloud ai-platform jobs describe $JOB_NAME` command"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9915e5549302"
      },
      "outputs": [],
      "source": [
        "!gcloud ai-platform jobs describe $JOB_NAME"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6e5947c13a83"
      },
      "source": [
        "You can stream logs using `gcloud ai-platform jobs stream-logs $JOB_NAME`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aaf79b1f19f9"
      },
      "outputs": [],
      "source": [
        "!gcloud ai-platform jobs stream-logs $JOB_NAME"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e636b52e5913"
      },
      "source": [
        "## Cleaning up Notebook Environment\n",
        "\n",
        "After you are done experimenting, you can either [STOP](https://cloud.google.com/ai-platform/notebooks/docs/shut-down) or DELETE the AI Notebook instance to prevent any  charges. If you want to save your work, you can choose to stop the instance instead.\n",
        "\n",
        "```\n",
        "# Stopping AI Platform Notebook instance\n",
        "gcloud notebooks instances stop example-instance --location=us-central1-a\n",
        "\n",
        "\n",
        "# Deleting AI Platform Notebook instance\n",
        "gcloud notebooks instances delete example-instance --location=us-central1-a\n",
        "```"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "n9qywopnIrJH",
        "545PP3o8IrJV",
        "7k8ge1L1IrJk"
      ],
      "name": "pytorch-text-classification-caip-training.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
